// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/core/plan_node.h>
#include <pollux/vector/complex_vector.h>

namespace kumo::pollux::exec {
    class Aggregate;

    /// Information needed to evaluate an aggregate function.
    struct AggregateInfo {
        /// Instance of the Aggregate class.
        std::unique_ptr<Aggregate> function;

        /// Indices of the input columns in the input RowVector.
        std::vector<column_index_t> inputs;

        /// Optional constant inputs. The size of this vector matches the size of
        /// 'inputs'. Non-constant inputs have null entries.
        std::vector<VectorPtr> constantInputs;

        /// Optional index of an input boolean column that should be used as a mask.
        std::optional<column_index_t> mask;

        /// Optional list of input columns that should be used to sort input rows
        /// before aggregating. Thes column may or may not overlap with 'inputs'.
        std::vector<column_index_t> sortingKeys;

        /// Optional list of sorting orders that goes with 'sortingKeys'.
        std::vector<core::SortOrder> sortingOrders;

        /// Boolean indicating whether inputs must be de-duplicated before
        /// aggregating.
        bool distinct{false};

        /// Index of the result column in the output RowVector.
        column_index_t output{0};

        /// Type of intermediate results. Used for spilling.
        TypePtr intermediateType;
    };

    class OperatorCtx;

    /// Translate an AggregationNode to a list of AggregationInfo, which could be
    /// a hash aggregation plan node or a streaming aggregation plan node. Ignore
    /// sorting properties if aggregate function is not sensitive to the order of
    /// inputs.
    ///
    /// @param aggregationNode Plan node of this aggregation.
    /// @param operatorCtx Operator context.
    /// @param numKeys Number of grouping keys.
    /// @param expressionEvaluator An Expression evaluator. It is used by an
    /// aggregate operator to compile and eval lambda expression. It should be
    /// initiated/assigned for at most one time.
    /// @param isStreaming Indicate whether this aggregation is streaming or not.
    /// Pass true if the aggregate operator is a StreamingAggregation and false if
    /// the aggregate operator is a HashAggregation. This parameter will be
    /// removed after sorted, distinct aggregation, and lambda functions support
    /// are added to StreamingAggregation.
    /// @return List of AggregationInfo.
    std::vector<AggregateInfo> toAggregateInfo(
        const core::AggregationNode &aggregationNode,
        const OperatorCtx &operatorCtx,
        uint32_t numKeys,
        std::shared_ptr<core::ExpressionEvaluator> &expressionEvaluator,
        bool isStreaming = false);

    /// Extract index of the 'mask' column for each aggregation from aggregations.
    /// Aggregations without masks use std::nullopt.
    std::vector<std::optional<column_index_t> > extractMaskChannels(
        const std::vector<AggregateInfo> &aggregates);
} // namespace kumo::pollux::exec
